import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test the update of beliefs on world states
according to the ouctome of the last action
'''

class TheCupWasPoisonedTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds=1
        self.world = World()
        self.gertrude = Agent('Gertrude')
        self.world.addAgent(self.gertrude)
        self.hamlet = Agent('Hamlet')
        self.world.addAgent(self.hamlet)
        self.agents = [self.gertrude,self.hamlet]
        self.world.setOrder(['Gertrude'])

    def addStates(self):
        #define states
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        self.world.defineState(None,'poisoned-cup',bool)
        self.world.setState(None,'poisoned-cup',True)
        
        for agent in self.agents:
            self.world.defineState(agent.name,'dead',bool)
            agent.setState('dead',False)
         
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))
        
    def addActions(self):
        #add actions and appropriate dynamics
        action = self.gertrude.addAction({'verb': 'drink'})
        tree = makeTree({'if': trueRow(stateKey(None,'poisoned-cup')),
                         True: setTrueMatrix(stateKey(self.gertrude.name,'dead')),
                         False: noChangeMatrix(stateKey(self.gertrude.name,'dead'))})
        self.world.setDynamics(stateKey(self.gertrude.name,'dead'),action,tree)
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree)
              
    def setModelAndRewards(self):
        #set models, rewards, and beliefs
        self.world.setModel('Hamlet', True)

        key = stateKey(None,'poisoned-cup')
        for agent in self.agents:
            agent.setBelief(key,Distribution({0:1,1:0.}),True)        

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelAndRewards()

    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        self.world.printModels()
        for t in range(self.maxRounds + 1):
            if t == 0:
                #assert disbelief in poison
                beliefs = self.hamlet.getAttribute('beliefs', True).domain()
                self.assertEqual(beliefs[0]['poisoned-cup'], 0.0)
                self.assertEqual(beliefs[1]['poisoned-cup'], 1.0)
            else:
                #assert belief in poison
                beliefs = self.hamlet.getAttribute('beliefs', True).domain()
                self.assertEqual(beliefs[0]['poisoned-cup'], 1.0)
                
            self.world.explain(self.world.step(),level=1)
            self.world.printState(beliefs=True)
            self.world.printModels()
            if self.world.terminated():
                break 

if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(TheCupWasPoisonedTestCase):
       if method.startswith("test"):
          suite.addTest(TheCupWasPoisonedTestCase(method))
    unittest.TextTestRunner().run(suite)
            